#!/usr/bin/env python3
import threading
import cmd
import time
import argparse
from cluster_fault_server import create_server, ClusterFaultServer

from motor.utils.logger import get_logger

logger = get_logger(__name__)


class FaultCLI(cmd.Cmd):
    intro = """
    集群故障模拟器CLI
    输入 'help' 或 '?' 查看可用命令
    输入 'quit' 退出
    """
    prompt = "fault> "

    def __init__(self, fault_server: ClusterFaultServer):
        super().__init__()
        self.fault_server = fault_server

    def do_generate_fault(self, arg):
        """生成故障消息: generate_fault <node_ip> <fault_level> [fault_codes...]
        例如: generate_fault 192.168.1.100 unhealthy 0x1234 0x5678
        """
        args = arg.split()
        if len(args) < 2:
            print("用法: generate_fault <node_ip> <fault_level> [fault_codes...]")
            print("例如: generate_fault 192.168.1.100 unhealthy 0x1234 0x5678")
            return

        node_ip = args[0]
        fault_level = args[1]
        fault_codes = args[2:] if len(args) > 2 else []

        # 验证fault_level
        valid_levels = ["healthy", "sub_healthy", "unhealthy"]
        if fault_level not in valid_levels:
            print(f"无效的故障等级。有效值: {valid_levels}")
            return

        try:
            self.fault_server.add_fault_message(
                node_ip=node_ip,
                fault_level=fault_level,
                fault_codes=fault_codes
            )
            print(f"已生成故障消息: 节点={node_ip}, 等级={fault_level}, 故障码={fault_codes}")
        except Exception as e:
            print(f"生成故障消息失败: {e}")

    def do_list_clients(self, arg):
        """列出已注册的客户端"""
        if not self.fault_server.clients:
            print("没有已注册的客户端")
        else:
            print("已注册的客户端:")
            for job_id, client in self.fault_server.clients.items():
                print(f"  - Job ID: {job_id}, Role: {client.role}")

    def do_simulate_scenario(self, arg):
        """模拟故障场景: simulate_scenario <scenario_name>
        可用场景: node_failure, device_failure, recovery
        """
        scenarios = {
            "node_failure": self._scenario_node_failure,
            "device_failure": self._scenario_device_failure,
            "recovery": self._scenario_recovery
        }

        if not arg:
            print("可用场景: node_failure, device_failure, recovery")
            return

        scenario = arg.strip()
        if scenario in scenarios:
            print(f"开始模拟场景: {scenario}")
            scenarios[scenario]()
        else:
            print(f"未知场景: {scenario}")
            print("可用场景: node_failure, device_failure, recovery")

    def _scenario_node_failure(self):
        """模拟节点故障场景"""
        print("模拟节点故障...")
        self.fault_server.add_fault_message(
            node_ip="192.168.1.100",
            fault_level="unhealthy",
            fault_codes=["0xDEAD", "0xBEEF"]
        )
        time.sleep(2)
        self.fault_server.add_fault_message(
            node_ip="192.168.1.101",
            fault_level="unhealthy",
            fault_codes=["0xCAFE"]
        )

    def _scenario_device_failure(self):
        """模拟设备故障场景"""
        print("模拟设备故障...")
        self.fault_server.add_fault_message(
            node_ip="192.168.1.100",
            fault_level="sub_healthy",
            fault_codes=["0x1234"],
            device_type="npu",
            device_id="0"
        )
        time.sleep(1)
        self.fault_server.add_fault_message(
            node_ip="192.168.1.100",
            fault_level="unhealthy",
            fault_codes=["0x5678"],
            device_type="npu",
            device_id="1"
        )

    def _scenario_recovery(self):
        """模拟恢复场景"""
        print("模拟恢复...")
        self.fault_server.add_fault_message(
            node_ip="192.168.1.100",
            fault_level="healthy",
            fault_codes=[]
        )

    def do_quit(self, arg):
        """退出CLI"""
        print("退出故障模拟器...")
        return True

    def do_help(self, arg):
        """显示帮助信息"""
        print("""
可用命令:
  generate_fault <node_ip> <fault_level> [fault_codes...] - 生成自定义故障消息
    故障等级: healthy, sub_healthy, unhealthy
    示例: generate_fault 192.168.1.100 unhealthy 0x1234 0x5678

  list_clients - 列出已注册的客户端

  simulate_scenario <scenario> - 运行预定义故障场景
    可用场景:
      node_failure - 模拟节点故障
      device_failure - 模拟设备故障
      recovery - 模拟恢复

  quit - 退出程序

输入 'help <command>' 获取特定命令的帮助
        """)


def main():
    parser = argparse.ArgumentParser(description="集群故障模拟器")
    parser.add_argument("--host", default="localhost", help="服务器主机地址")
    parser.add_argument("--port", type=int, default=5005, help="服务器端口")
    args = parser.parse_args()

    print("启动集群故障模拟器...")
    print(f"服务器地址: {args.host}:{args.port}")

    # 创建并启动服务器
    server, fault_server, server_address = create_server(args.host, args.port)
    server.start()

    print(f"服务器启动完成: {server_address}")
    print("启动CLI界面...")

    try:
        # 启动CLI
        cli = FaultCLI(fault_server)
        cli.cmdloop()
    except KeyboardInterrupt:
        print("\n收到中断信号，正在关闭...")
    finally:
        print("正在停止服务器...")
        fault_server.running = False
        server.stop(0)
        print("服务器已停止")


if __name__ == "__main__":
    main()
